// Copyright (C) 2022-2025 Exaloop Inc. <https://exaloop.io>

#include "const_prop.h"

#include "codon/cir/analyze/dataflow/reaching.h"
#include "codon/cir/analyze/module/global_vars.h"
#include "codon/cir/util/cloning.h"
#include "codon/cir/util/irtools.h"

namespace codon {
namespace ir {
namespace transform {
namespace folding {
namespace {
bool okConst(const Value *v) {
  return v && (isA<IntConst>(v) || isA<FloatConst>(v) || isA<BoolConst>(v));
}
} // namespace

const std::string ConstPropPass::KEY = "core-folding-const-prop";

void ConstPropPass::handle(VarValue *v) {
  auto *M = v->getModule();
  auto *var = v->getVar();
  Value *replacement;

  if (var->isGlobal()) {
    auto *r = getAnalysisResult<analyze::module::GlobalVarsResult>(globalVarsKey);
    if (!r)
      return;

    auto it = r->assignments.find(var->getId());
    if (it == r->assignments.end())
      return;

    auto *constDef = M->getValue(it->second);
    if (!okConst(constDef))
      return;

    util::CloneVisitor cv(M);
    replacement = cv.clone(constDef);
  } else {
    auto *r = getAnalysisResult<analyze::dataflow::RDResult>(reachingDefKey);
    if (!r)
      return;
    auto *c = r->cfgResult;

    auto it = r->results.find(getParentFunc()->getId());
    if (it == r->results.end())
      return;

    auto *rd = it->second.get();
    auto reaching = rd->getReachingDefinitions(var, v);

    if (reaching.size() != 1 || !reaching[0].known())
      return;

    auto *constDef = reaching[0].assignee;
    if (!okConst(constDef))
      return;

    util::CloneVisitor cv(M);
    replacement = cv.clone(constDef);
  }

  v->replaceAll(replacement);
}

void ConstPropPass::handle(ExtractInstr *v) {
  // Propagate constant tuples

  if (!isA<VarValue>(v->getVal()) || !isA<types::RecordType>(v->getVal()->getType()))
    return;

  auto *var = cast<VarValue>(v->getVal())->getVar();
  if (var->isGlobal())
    return;

  auto *r = getAnalysisResult<analyze::dataflow::RDResult>(reachingDefKey);
  if (!r)
    return;

  auto *c = r->cfgResult;
  auto it = r->results.find(getParentFunc()->getId());
  if (it == r->results.end())
    return;

  auto *rd = it->second.get();
  auto reaching = rd->getReachingDefinitions(var, v);
  if (reaching.size() != 1 || !reaching[0].known())
    return;

  auto *call = cast<CallInstr>(reaching[0].assignee);
  if (!call)
    return;

  auto *func = util::getFunc(call->getCallee());
  if (!func || func->getUnmangledName() != Module::NEW_MAGIC_NAME)
    return;

  auto *tuple = cast<types::RecordType>(func->getParentType());
  if (!tuple || tuple->getName() != "Tuple")
    return;

  auto idx =
      cast<types::RecordType>(v->getVal()->getType())->getMemberIndex(v->getField());
  if (idx < 0 || idx >= call->numArgs() || !okConst(*(call->begin() + idx)))
    return;

  util::CloneVisitor cv(v->getModule());
  auto *replacement = cv.clone(*(call->begin() + idx));
  v->replaceAll(replacement);
}

} // namespace folding
} // namespace transform
} // namespace ir
} // namespace codon
